"""
postprocess
"""
import mindspore.nn as nn
import mindspore.numpy as msnp
from mindspore import ops


class PostProcess2DMinMax(nn.Cell):
    """
    2D模型MinMax归一化后的后处理模块
    将模型预测结果转换为物理空间中的真实涡粘性值
    """
    def __init__(self, reynolds, dis, sca_max, sca_min):
        """
        初始化后处理模块
        
        Args:
            reynolds: 雷诺数
            dis: 到壁面的距离
            sca_max: 归一化最大值
            sca_min: 归一化最小值
        """
        super(PostProcess2DMinMax, self).__init__()
        self.mul = ops.Mul()                # 乘法操作
        self.add = ops.Add()                # 加法操作
        self.pow = ops.Pow()                # 幂次操作
        self.low_value = 0.0                # 预测值下限（物理约束）
        self.up_value = 1.0                 # 预测值上限（物理约束）
        self.sca_max = sca_max              # 归一化最大值
        self.sca_min = sca_min              # 归一化最小值
        self.reynolds = reynolds            # 雷诺数
        self.dis = dis                      # 到壁面的距离

    def construct(self, pred):
        """
        后处理流程
        
        Args:
            pred: 模型预测的归一化涡粘性值
            
        Returns:
            Tensor: 物理空间中的真实涡粘性值
        """
        # 1、限制预测值范围，确保符合物理意义（涡粘性非负）
        pred = msnp.clip(pred, self.low_value, self.up_value)
        
        # 2、反归一化：将预测值从[0,1]范围映射回原始数据范围
        pred = pred[:, 0]  # 提取预测值
        pred = self.add(self.mul(pred, (self.sca_max - self.sca_min)), self.sca_min)
        
        # 3、逆尺度变换：恢复物理空间中的真实涡粘性值
        # 距离相关的尺度因子（壁面附近涡粘性变化剧烈）
        trans = 1.0 / self.pow(self.dis, 0.6)
        # 雷诺数尺度调整
        reynolds = self.reynolds * 1e-6
        # 最终物理涡粘性值
        pred = pred / trans * reynolds
        
        return pred


class PostProcess2DStd(nn.Cell):
    """
    2D模型标准差归一化后的后处理模块
    将模型预测结果转换为物理空间中的真实涡粘性值
    """
    def __init__(self, reynolds, dis, df_mean, df_std):
        """
        初始化后处理模块
        
        Args:
            reynolds: 雷诺数
            dis: 到壁面的距离
            df_mean: 训练数据均值
            df_std: 训练数据标准差
        """
        super(PostProcess2DStd, self).__init__()
        self.mul = ops.Mul()                # 乘法操作
        self.add = ops.Add()                # 加法操作
        self.pow = ops.Pow()                # 幂次操作
        self.low_value = 0.0                # 预测值下限
        self.up_value = 1.0                 # 预测值上限
        self.df_mean = df_mean              # 训练数据均值
        self.df_std = df_std                # 训练数据标准差
        self.reynolds = reynolds            # 雷诺数
        self.dis = dis                      # 到壁面的距离

    def construct(self, pred):
        """
        后处理流程
        
        Args:
            pred: 模型预测的标准化涡粘性值
            
        Returns:
            Tensor: 物理空间中的真实涡粘性值
        """
        # 1、反归一化：将标准化预测值恢复到原始数据分布
        pred = pred[:, 0]  # 提取预测值
        pred = self.add(self.mul(pred, self.df_std), self.df_mean)
        
        # 2、逆尺度变换：恢复物理空间中的真实涡粘性值
        # 距离相关的尺度因子
        trans = 1.0 / self.pow(self.dis, 0.6)
        # 雷诺数尺度调整
        reynolds = self.reynolds * 1e-6
        # 最终物理涡粘性值
        pred = pred / trans * reynolds
        
        return pred


class PostProcess3DMinMax(nn.Cell):
    """
    3D模型MinMax归一化后的后处理模块
    将模型预测结果转换为物理空间中的真实涡粘性值
    """
    def __init__(self, reynolds, sca_max, sca_min):
        """
        初始化后处理模块
        
        Args:
            reynolds: 雷诺数
            sca_max: 归一化最大值
            sca_min: 归一化最小值
        """
        super(PostProcess3DMinMax, self).__init__()
        self.mul = ops.Mul()                # 乘法操作
        self.add = ops.Add()                # 加法操作
        self.low_value = 0.0                # 预测值下限
        self.up_value = 1.0                 # 预测值上限
        self.sca_max = sca_max              # 归一化最大值
        self.sca_min = sca_min              # 归一化最小值
        self.reynolds = reynolds            # 雷诺数

    def construct(self, pred):
        """
        后处理流程
        
        Args:
            pred: 模型预测的归一化涡粘性值
            
        Returns:
            Tensor: 物理空间中的真实涡粘性值
        """
        # 1、限制预测值范围，确保符合物理意义
        pred = msnp.clip(pred, self.low_value, self.up_value)
        
        # 2、反归一化：将预测值从[0,1]范围映射回原始数据范围
        pred = pred[:, 0]  # 提取预测值
        pred = self.add(self.mul(pred, (self.sca_max-self.sca_min)), self.sca_min)
        
        # 3、逆尺度变换：恢复物理空间中的真实涡粘性值
        # 雷诺数尺度调整
        reynolds = self.reynolds * 1e-6
        # 最终物理涡粘性值
        pred = pred * reynolds
        
        return pred